#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/4/11 11:01
# @Author  : liutianwei
# @File    : snr_test.py
# @Software: PyCharm

import os.path
import re
import sys
sys.path.insert(0, './build')
import snr_calculate


class snrCalculate:
    def __init__(self, **kwargs):
        self.__voicetype = "wave"
        self.__output_path = ""
        self.__sp_size = 30
        self.__sil_size = 30
        self.__frame_size = 320
        self.parse_case(kwargs)
        return

    def parse_case(self, kwargs):
        keys = kwargs.keys()
        if "frame_size" in keys:
            self.__frame_size = kwargs["frame_size"]
        if "output_path" in keys:
            self.__output_path = os.path.join(kwargs["output_path"], "snr.txt")
        if "sp_size" in keys:
            self.__sp_size = kwargs["sp_size"]
        if "sil_size" in keys:
            self.__sil_size = kwargs["sil_size"]
        return

    def calculate_snr_time(self, wav_path, voice_time, sample_rate):
        lbl_file = []
        # 计算音频的结束时间
        channelnum = int(self.__frame_size / 320)
        audio_length = os.path.getsize(wav_path)
        with open(wav_path, 'rb') as fp:
            # 获取文件头信息
            header = fp.read(100)
            if header[:4] == b"RIFF":
                # 如果是wav文件，则获取头信息中的音频长度
                pos = re.search(b'data', header).end()
                audio_length = int.from_bytes(header[pos:pos + 4], byteorder='little')
                # 将文件指针移动到音频数据的起始位置
                fp.seek(pos + 4, 0)
            else:
                self.__voicetype = 'pcm'
                fp.seek(0, 0)

        # 实际pcm音频流数据
        remaining_bytes = audio_length
        filesize = remaining_bytes
        voice_end_time = 1000 * filesize / (2.0 * sample_rate * channelnum)
        voice_end_time = round(voice_end_time/1000, 2)
        time_node = ['0.00']
        # 判断传入的lbl是列表还是文件
        if voice_time is not None:
            if type(voice_time) is not list:
                with open(voice_time) as f:
                    for line in f:
                        f_list = line.split(' ')
                        f_list = [item.strip() for item in f_list]
                        for f_item in f_list:
                            if f_item == "":
                                continue
                            time_node.append(f_item)
                    time_node.append(str(voice_end_time))
            else:
                # 拆分字符串转成lbl列表
                for time in voice_time:
                    time_temp = time.lstrip('[')
                    time_temp = time_temp.rstrip(']')
                    time_node.append(time_temp)
                time_node.append(str(voice_end_time))
        else:
            time_node.append(str(voice_end_time))
        time_mark = 'N'
        for i in range(len(time_node)):
            if i == len(time_node) - 1:
                break
            lbl_file.append([time_node[i], time_node[i+1], time_mark])
            if time_mark == 'N':
                time_mark = 'S'
            else:
                time_mark = 'N'
        return lbl_file

    def snr_calculate_model(self, wav_path, lbl_path, sample_rate):
        # 对输入的参数进行校验

        if wav_path == "":
            print("Not a wav file")
            return False
        if lbl_path == "":
            print("Not a lbl file")
            return False
        if self.__output_path == "":
            print("Not a output path")
            return False
        if sample_rate !=8000 and sample_rate != 16000 :
            print("sample_rate is round 8000 and 16000")
            return False
        lbl_file = self.calculate_snr_time(wav_path, lbl_path, sample_rate)
        snr_result = snr_calculate.snrparse(wav_path, lbl_file, self.__output_path, sample_rate,
                                            self.__voicetype, self.__sp_size, self.__sil_size)
        return snr_result

if __name__ == "__main__":
    test_snr_calculate = snrCalculate( sample_rate = 16000, sp_size =30, sil_size =30, output_path='/data1/liutianwei/mango', frame_size = 320)
    result = test_snr_calculate.snr_calculate_model("/data1/liutianwei/mango/讲一个绕口令_cmn.wav",'/data1/liutianwei/mango/lbl.txt')
    result1 = test_snr_calculate.snr_calculate_model("/data1/liutianwei/mango/讲一个绕口令_cmn.wav", '[[0.54,1.73]]')
    result2 = test_snr_calculate.snr_calculate_model("/data1/liutianwei/mango/讲一个绕口令_cmn.wav", '/data1/liutianwei/mango/lbl.lbl')
    print(result)
    print(result1)
    print(result2)


